# TODO 多进程模式，希望解决延迟高
import numpy as np
import multiprocessing as mp
from multiprocessing import shared_memory


def process_A(shape, data_ready_event, result_ready_event):
    # 创建共享内存块用于发送数据m和接收结果k


    # 将共享内存包装为np.array
    m = np.ndarray(shape, dtype=np.float32, buffer=shm_m.buf)
    k = np.ndarray(shape, dtype=np.float32, buffer=shm_k.buf)

    # 模拟图形渲染生成数据m
    print("A generating data m")
    m[:] = np.array([1,2,3,4,5,6]).astype(np.float32)

    # 通知B数据已就绪
    print("A data ready")
    data_ready_event.set() # 得到mocap位置后通知control

    # 等待B计算完成
    print("A waiting for B to finish")
    result_ready_event.wait() # MUJOCO里面先把mocap发给control，然后再wait，等待返回的ik_qstar
    print("A received result:", k)

    # 清理共享内存
    shm_m.close()
    shm_m.unlink()
    shm_k.close()
    shm_k.unlink()

def process_B(shape, data_ready_event, result_ready_event):
    # 连接到A创建的共享内存
    

    print("B connected to A")
    # 包装为np.array
    for i in range(10):
        m = np.ndarray(shape, dtype=np.float32, buffer=shm_m.buf)
        k = np.ndarray(shape, dtype=np.float32, buffer=shm_k.buf)
        print('get m')
        m += 1
        print(m)
    # 等待A的数据就绪
    print("B waiting for A to finish")
    data_ready_event.wait() # 等待mujoco给出mocap位置

    # 实时计算
    k[:] = m+1  # 此处替换为实际计算逻辑

    # 通知A结果已就绪
    result_ready_event.set() # 求出ik_qstar后就通知mujoco

    # 清理共享内存
    shm_m.close()
    shm_k.close()

if __name__ == "__main__":
    shape = (6, )  # 示例数据维度
    print("A creating shared memory")
    shm_m = shared_memory.SharedMemory(create=True, size=np.prod(shape)*4)  # float32
    shm_k = shared_memory.SharedMemory(create=True, size=np.prod(shape)*4)
    global shm_m_name, shm_k_name
    shm_m_name = shm_m.name
    shm_k_name = shm_k.name
    # 创建同步事件
    data_ready_event = mp.Event()
    result_ready_event = mp.Event()



    # process_B
    p_b = mp.Process(target=process_B, args=(shape, data_ready_event, result_ready_event))
    p_b.start()

    # process_A
    p_a = mp.Process(target=process_A, args=(shape, data_ready_event, result_ready_event))
    p_a.start()



    p_a.join()
    p_b.join()